#!/usr/bin/env python
# -*- encoding: UTF-8 -*-
from __future__ import print_function

"""
正逆运动学的符号分析库
"""

from sympy import *
from nao_defines import *

# DH 参数矩阵
def DH(a, alpha, d, theta):
    from sympy import cos,sin
    Taf = [[cos(theta),            -sin(theta),            0,              a],
    [sin(theta)*cos(alpha), cos(theta)*cos(alpha),  -sin(alpha),    -sin(alpha)*d],
    [sin(theta)*sin(alpha), cos(theta)*sin(alpha),  cos(alpha),     cos(alpha)*d],
    [0,                     0,                      0,              1]]
    return Matrix(Taf)

def RotXYZMatrix(xAngle,yAngle,zAngle):
    from symrobot import rotx,roty,rotz,r2t
    Rx = rotx(xAngle)
    Ry = roty(yAngle)
    Rz = rotz(zAngle)
    R = Rx*Ry*Rz
    return r2t(R)

def RotZYXMatrix(zAngle,yAngle,xAngle):
    from symrobot import rotx,roty,rotz,r2t
    Rx = rotx(xAngle)
    Ry = roty(yAngle)
    Rz = rotz(zAngle)
    R = Rz*Ry*Rx
    return r2t(R)

# 单例修饰器
def singleton(cls, *args, **kw):
    instances = {}
    def _singleton():
        if cls not in instances:
            instances[cls] = cls(*args, **kw)
        return instances[cls]
    return _singleton

# 获取符号Nao机器人参数的表示
@singleton
class Get_symbols(object):
    def __init__(self):
        thetas = [(Symbol("theta_%d" % n), n)  for n in range(1,6)]
        for theta, n in thetas:
            self.__dict__["theta%d" % n] = theta
        self.__dict__["theta%d" % n] = Symbol("theta")

    def __getattr__(self, item):
        if item in self.__dict__:
            return self.__dict__[item]
        else:
            self.__dict__[item] = Symbol(item)
            return self.__dict__[item]

    def tihuan(self,S):
        """ 替换规则 """
        return S.subs([(self.UpperArmLength,Symbol("l_1")),
                       (self.ElbowOffsetY,Symbol("l_2")),
                       (self.ShoulderOffsetY,Symbol("l_3")),
                       (self.ShoulderOffsetZ,Symbol("l_4")),
                       (self.HandOffsetX,Symbol("l_5")),
                       (self.HandOffsetZ,Symbol("l_6")),
                       (self.LowerArmLength,Symbol("l_7"))
                        ])
    def evalf(self,S):
        return S.subs([(self.UpperArmLength,UpperArmLength),
                       (self.ElbowOffsetY,ElbowOffsetY),
                       (self.ShoulderOffsetY,ShoulderOffsetY),
                       (self.ShoulderOffsetZ,ShoulderOffsetZ),
                       (self.HandOffsetX,HandOffsetX),
                       (self.HandOffsetZ,HandOffsetZ),
                       (self.LowerArmLength,LowerArmLength)
                        ])

class LeftHand(object):
    def __init__(self):
        sb = Get_symbols()
        self.sb = sb
        self.Base = Matrix.eye(4)
        self.Base[1,3] = sb.ShoulderOffsetY
        self.Base[2,3] = sb.ShoulderOffsetZ
        self.T1 = DH(0,-pi/2,0,sb.theta1)        # LShoulderPitch
        self.T2 = DH(0,pi/2,0,sb.theta2+pi/2)    # LShoulderRoll
        self.T3 = DH(sb.ElbowOffsetY,pi/2,sb.UpperArmLength,sb.theta3) # LElbowYaw
        self.T4 = DH(0,-pi/2,0,sb.theta4)        # LElbowRoll
        self.T5 = DH(0,pi/2,0,sb.theta5)         # LWristYaw
        self.Tool = Matrix.eye(4)
        self.Tool[0,3] = sb.LowerArmLength + sb.HandOffsetX
        self.Tool[2,3] = -sb.HandOffsetZ
        self.Tool = RotXYZMatrix(-pi/2,0,-pi/2)*self.Tool
        self.All = self.Base*self.T1*self.T2*self.T3*self.T4*self.T5*self.Tool

    def forward(self,th1,th2,th3,th4,th5):
        All = self.sb.evalf(self.All)
        All = All.subs([(self.sb.theta1,th1),
                  (self.sb.theta2,th2),
                  (self.sb.theta3,th3),
                  (self.sb.theta4,th4),
                  (self.sb.theta5,th5)
                  ])
        return All.evalf()

class RightHand(object):
    def __init__(self):
        sb = Get_symbols()
        self.sb = sb
        self.Base = Matrix.eye(4)
        self.Base[1,3] = -sb.ShoulderOffsetY # 与LArm不同
        self.Base[2,3] = sb.ShoulderOffsetZ
        self.T1 = DH(0,-pi/2,0,sb.theta1)       # RShoulderPitch
        self.T2 = DH(0,pi/2,0,sb.theta2+pi/2)   # RShoulderRoll
        self.T3 = DH(-sb.ElbowOffsetY,pi/2,sb.UpperArmLength,sb.theta3) # RElbowYaw 与 LArm 不同
        self.T4 = DH(0,-pi/2,0,sb.theta4)       # RElbowRoll
        self.T5 = DH(0,pi/2,0,sb.theta5)        # RWristYaw
        self.Tool = Matrix.eye(4)
        self.Tool[0,3] = sb.LowerArmLength + sb.HandOffsetX
        self.Tool[2,3] = -sb.HandOffsetZ
        self.Tool = RotXYZMatrix(-pi/2,0,-pi/2)*self.Tool
        self.All = self.Base*self.T1*self.T2*self.T3*self.T4*self.T5*self.Tool

    def forward(self,th1,th2,th3,th4,th5):
        All = self.sb.evalf(self.All)
        All = All.subs([(self.sb.theta1,th1),
                  (self.sb.theta2,th2),
                  (self.sb.theta3,th3),
                  (self.sb.theta4,th4),
                  (self.sb.theta5,th5)
                  ])
        return All.evalf()

class LeftLeg(object):
    def __init__(self):
        sb = Get_symbols()
        self.sb = sb
        self.Base = Matrix.eye(4)
        self.Base[1,3] = sb.HipOffsetY
        self.Base[2,3] = -sb.HipOffsetZ
        self.T1 = DH(0,-3*pi/4,0,sb.theta1-pi/2)  # LHipYawPitch
        self.T2 = DH(0,-pi/2,0,sb.theta2+pi/4)    # LHipRoll
        self.T3 = DH(0,pi/2,0,sb.theta3)          # LHipPitch
        self.T4 = DH(-sb.ThighLength,0,0,sb.theta4)  # LKneePitch
        self.T5 = DH(-sb.TibiaLength,0,0,sb.theta5)  # LAnklePitch
        self.T6 = DH(0,-pi/2,0,sb.theta6)         # LAnkleRoll
        self.Tool = Matrix.eye(4)
        self.Tool[2,3] = -sb.FootHeight
        self.Tool = RotZYXMatrix(pi,-pi/2,0) * self.Tool
        self.All = self.Base*self.T1*self.T2*self.T3*self.T4*self.T5*self.T6*self.Tool

    def forward(self,th1,th2,th3,th4,th5,th6):
        All = self.sb.evalf(self.All)
        All = All.subs([(self.sb.theta1,th1),
                  (self.sb.theta2,th2),
                  (self.sb.theta3,th3),
                  (self.sb.theta4,th4),
                  (self.sb.theta5,th5),
                  (self.sb.theta6,th6)
                  ])
        return All.evalf()

class RightLeg(object):
    def __init__(self):
        sb = Get_symbols()
        self.sb = sb
        self.Base = Matrix.eye(4)
        self.Base[1,3] = -sb.HipOffsetY # 与 LeftLeg 不同
        self.Base[0,3] = -sb.HipOffsetZ
        self.T1 = DH(0,-pi/4,0,sb.theta1-pi/2)   # RHipYawPitch 与 LeftLeg 不同
        self.T2 = DH(0,-pi/2,0,sb.theta2-pi/4)   # RHipRoll 与 LeftLeg 不同
        self.T3 = DH(0,pi/2,0,sb.theta3)         # RHipPitch
        self.T4 = DH(-sb.ThighLength,0,0,sb.theta4) # RKneePitch
        self.T5 = DH(-sb.TibiaLength,0,0,sb.theta5) # RAnklePitch
        self.T6 = DH(0,-pi/2,0,sb.theta6)        # RAnkleRoll
        self.Tool = Matrix.eye(4)
        self.Tool[2,3] = -sb.FootHeight
        self.Tool = RotZYXMatrix(pi,-pi/2,0) * self.Tool
        self.All = self.Base*self.T1*self.T2*self.T3*self.T4*self.T5*self.T6*self.Tool
        self.Jacobian = self.All.jacobian([sb.theta1,sb.theta2,sb.theta3,sb.theta4,sb.theta5,sb.theta6])

    def forward(self,th1,th2,th3,th4,th5,th6):
        All = self.sb.evalf(self.All)
        All = All.subs([(self.sb.theta1,th1),
                  (self.sb.theta2,th2),
                  (self.sb.theta3,th3),
                  (self.sb.theta4,th4),
                  (self.sb.theta5,th5),
                  (self.sb.theta6,th6)
                  ])
        return All.evalf()

    def JacobianV(self,th1,th2,th3,th4,th5,th6):
        Jacobian = self.sb.evalf(self.Jacobian)
        Jacobian = Jacobian.subs([(self.sb.theta1,th1),
                  (self.sb.theta2,th2),
                  (self.sb.theta3,th3),
                  (self.sb.theta4,th4),
                  (self.sb.theta5,th5),
                  (self.sb.theta6,th6)
                  ])
        return Jacobian.evalf()

class Head(object):
    def __init__(self):
        sb = Get_symbols()
        self.sb = sb
        self.Base = Matrix.eye(4)
        self.Base[2,3] = sb.NeckOffsetZ
        self.T1 = DH(0,0,0,sb.theta1)            # HeadYaw
        self.T2 = DH(0,-pi/2,0,sb.theta2-pi/2)   # HeadPitch
        ## ???
        self.ToolDownCamera = Matrix.eye(4)
        self.ToolDownCamera[0,3] = sb.CameraBotomX
        self.ToolDownCamera[2,3] = sb.CameraBotomZ
        self.ToolDownCamera = RotXYZMatrix(pi/2,pi/2,0) * self.ToolDownCamera
        self.AllDownCamera = self.Base*self.T1*self.T2*self.ToolDownCamera

@singleton
class Kinematics(object):
    def __init__(self):
        sb = Get_symbols()
        self.sb = sb
        self.Head = Head()
        self.LeftHand = LeftHand()
        self.RightHand = RightHand()
        self.LeftLeg = LeftLeg()
        self.RightLeg = RightLeg()

#######################################################################################################################

# 反对称矩阵
def ss(b):
    return  Matrix([
        [0, -b[2],  b[1]],
        [b[2],  0,  -b[0]],
        [-b[1], b[0],   0]])

# 旋转 Rodrigues 矩阵
def get_Rodrigues(r=None):
    theta = symbols("theta")
    if r is None:
        rx,ry,rz = symbols("r_x r_y r_z")
        r = Matrix([rx,ry,rz])
    R = cos(theta)*Matrix.eye(3)+(1-cos(theta))*r*r.transpose()+sin(theta)*ss(r)
    return R

def get_inv_Rodrigues(r=None):
    theta = symbols("theta")
    if r is None:
        rx,ry,rz = symbols("r_x r_y r_z")
        r = Matrix([rx,ry,rz])
    R = cos(theta)*Matrix.eye(3)+(1-cos(theta))*r*r.transpose()+sin(theta)*ss(r)
    return R

def eval_Rodrigues(w):
    theta,rx,ry,rz = symbols("theta r_x r_y r_z")
    real_theta = np.linalg.norm(w)
    w = w/real_theta
    R = get_Rodrigues()
    return R.subs([(theta,real_theta),(rx,w[0]),(ry,w[1]),(rz,w[2])]).evalf()

def get_Homg():
    T = []
    for i in range(3):
        T.append([])
        for j in range(4):
            T[i].append(Symbol("r_{%d,%d}" % (i+1,j+1)))
    T.append([0,0,0,1])
    return Matrix(T)
    
    
# 获取绕指定轴转动的后末端位姿符号矩阵
def get_Tend():
    rx,ry,rz = symbols("r_x r_y r_z")
    r = Matrix([rx,ry,rz])
    T = get_Homg()
    Trodrigues = Matrix.eye(4)
    R = get_Rodrigues(r)
    Trodrigues[:3,:3] = R
    T = T*Trodrigues
    return T

def get_Tend_k():
    k = symbols("k")
    rx,ry,rz = symbols("r_x r_y r_z")
    r = Matrix([rx,ry,rz])
    T = get_Homg()
    Trodrigues = Matrix.eye(4)
    R = get_Rodrigues(r)
    Trodrigues[:3,:3] = R
    T = T*Trodrigues
    p = r*k
    T[:3,3] += p
    return T